import cv2
import uvicorn
from fastapi import FastAPI, Response
from fastapi.responses import StreamingResponse


"""
pip install opencv-python
"""
# 创建一个FastAPI应用程序。
app = FastAPI()

# 设置RTSP URL。在这个例子中，我们将使用一个IP摄像头作为RTSP源。
rtsp_url = "rtsp://admin:Telchina2002@10.90.23.232:554/h264/ch1/main/av_stream?videoCodecType=H.264"

# 定义一个辅助函数，用于从RTSP源读取视频流。
def generate_frames():
    cap = cv2.VideoCapture(rtsp_url)
    while True:
        # 从视频流中读取帧。
        ret, frame = cap.read()
        if not ret:
            break

        # 将帧转换为JPEG格式。
        ret, buffer = cv2.imencode(".jpg", frame)

        # 将JPEG数据转换为字节字符串，并将其作为流响应返回。
        yield (b"--frame\r\n"
               b"Content-Type: image/jpeg\r\n\r\n" + buffer.tobytes() + b"\r\n")

    # 释放视频流。
    cap.release()

# 定义一个FastAPI路由，用于处理视频流请求。
@app.get("/video_feed")
async def video_feed():
    """将视频流作为流响应返回。"""
    return StreamingResponse(generate_frames(), media_type="multipart/x-mixed-replace;boundary=frame")

# 定义一个FastAPI路由，用于呈现网站页面。
@app.get("/")
async def root():
    """呈现网站页面。"""
    html = """
    <html>
        <head>
            <title>RTSP Video Streaming</title>
        </head>
        <body>
            <h1>RTSP Video Streaming</h1>
            <img src="/video_feed">
        </body>
    </html>
    """
    return Response(content=html, media_type="text/html")

if __name__ == '__main__':
    uvicorn.run(app=app)
